
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN"
  "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">

<html xmlns="http://www.w3.org/1999/xhtml" lang="en">
  <head>
    <meta http-equiv="X-UA-Compatible" content="IE=Edge" />
    <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
    <title>tensorflowonspark.TFSparkNode &#8212; TensorFlowOnSpark 1.3.3 documentation</title>
    <link rel="stylesheet" href="../../_static/classic.css" type="text/css" />
    <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
    <script type="text/javascript" id="documentation_options" data-url_root="../../" src="../../_static/documentation_options.js"></script>
    <script type="text/javascript" src="../../_static/jquery.js"></script>
    <script type="text/javascript" src="../../_static/underscore.js"></script>
    <script type="text/javascript" src="../../_static/doctools.js"></script>
    <link rel="index" title="Index" href="../../genindex.html" />
    <link rel="search" title="Search" href="../../search.html" /> 
  </head><body>
    <div class="related" role="navigation" aria-label="related navigation">
      <h3>Navigation</h3>
      <ul>
        <li class="right" style="margin-right: 10px">
          <a href="../../genindex.html" title="General Index"
             accesskey="I">index</a></li>
        <li class="right" >
          <a href="../../py-modindex.html" title="Python Module Index"
             >modules</a> |</li>
        <li class="nav-item nav-item-0"><a href="../../index.html">TensorFlowOnSpark 1.3.3 documentation</a> &#187;</li>
          <li class="nav-item nav-item-1"><a href="../index.html" accesskey="U">Module code</a> &#187;</li> 
      </ul>
    </div>  

    <div class="document">
      <div class="documentwrapper">
        <div class="bodywrapper">
          <div class="body" role="main">
            
  <h1>Source code for tensorflowonspark.TFSparkNode</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2017 Yahoo Inc.</span>
<span class="c1"># Licensed under the terms of the Apache 2.0 license.</span>
<span class="c1"># Please see LICENSE file in the project root for terms.</span>
<span class="sd">&quot;&quot;&quot;This module provides low-level functions for managing the TensorFlowOnSpark cluster.&quot;&quot;&quot;</span>

<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">absolute_import</span>
<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">division</span>
<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">nested_scopes</span>
<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">print_function</span>

<span class="kn">import</span> <span class="nn">json</span>
<span class="kn">import</span> <span class="nn">logging</span>
<span class="kn">import</span> <span class="nn">multiprocessing</span>
<span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">platform</span>
<span class="kn">import</span> <span class="nn">socket</span>
<span class="kn">import</span> <span class="nn">subprocess</span>
<span class="kn">import</span> <span class="nn">sys</span>
<span class="kn">import</span> <span class="nn">uuid</span>
<span class="kn">import</span> <span class="nn">time</span>
<span class="kn">import</span> <span class="nn">traceback</span>
<span class="kn">from</span> <span class="nn">threading</span> <span class="k">import</span> <span class="n">Thread</span>

<span class="kn">from</span> <span class="nn">.</span> <span class="k">import</span> <span class="n">TFManager</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="k">import</span> <span class="n">TFNode</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="k">import</span> <span class="n">gpu_info</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="k">import</span> <span class="n">marker</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="k">import</span> <span class="n">reservation</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="k">import</span> <span class="n">util</span>


<div class="viewcode-block" id="TFNodeContext"><a class="viewcode-back" href="../../tensorflowonspark.TFSparkNode.html#tensorflowonspark.TFSparkNode.TFNodeContext">[docs]</a><span class="k">class</span> <span class="nc">TFNodeContext</span><span class="p">:</span>
  <span class="sd">&quot;&quot;&quot;Encapsulates unique metadata for a TensorFlowOnSpark node/executor and provides methods to interact with Spark and HDFS.</span>

<span class="sd">  An instance of this object will be passed to the TensorFlow &quot;main&quot; function via the `ctx` argument.</span>
<span class="sd">  To simply the end-user API, this class now mirrors the functions of the TFNode module.</span>

<span class="sd">  Args:</span>
<span class="sd">    :executor_id: integer identifier for this executor, per ``nodeRDD = sc.parallelize(range(num_executors), num_executors).``</span>
<span class="sd">    :job_name: TensorFlow job name (e.g. &#39;ps&#39; or &#39;worker&#39;) of this TF node, per cluster_spec.</span>
<span class="sd">    :task_index: integer rank per job_name, e.g. &quot;worker:0&quot;, &quot;worker:1&quot;, &quot;ps:0&quot;.</span>
<span class="sd">    :cluster_spec: dictionary for constructing a tf.train.ClusterSpec.</span>
<span class="sd">    :defaultFS: string representation of default FileSystem, e.g. ``file://`` or ``hdfs://&lt;namenode&gt;:8020/``.</span>
<span class="sd">    :working_dir: the current working directory for local filesystems, or YARN containers.</span>
<span class="sd">    :mgr: TFManager instance for this Python worker.</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">executor_id</span><span class="p">,</span> <span class="n">job_name</span><span class="p">,</span> <span class="n">task_index</span><span class="p">,</span> <span class="n">cluster_spec</span><span class="p">,</span> <span class="n">defaultFS</span><span class="p">,</span> <span class="n">working_dir</span><span class="p">,</span> <span class="n">mgr</span><span class="p">):</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">worker_num</span> <span class="o">=</span> <span class="n">executor_id</span>       <span class="c1"># for backwards-compatibility</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">executor_id</span> <span class="o">=</span> <span class="n">executor_id</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">job_name</span> <span class="o">=</span> <span class="n">job_name</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">task_index</span> <span class="o">=</span> <span class="n">task_index</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">cluster_spec</span> <span class="o">=</span> <span class="n">cluster_spec</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">defaultFS</span> <span class="o">=</span> <span class="n">defaultFS</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">working_dir</span> <span class="o">=</span> <span class="n">working_dir</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">mgr</span> <span class="o">=</span> <span class="n">mgr</span>

<div class="viewcode-block" id="TFNodeContext.absolute_path"><a class="viewcode-back" href="../../tensorflowonspark.TFSparkNode.html#tensorflowonspark.TFSparkNode.TFNodeContext.absolute_path">[docs]</a>  <span class="k">def</span> <span class="nf">absolute_path</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">path</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Convenience function to access ``TFNode.hdfs_path`` directly from this object instance.&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">TFNode</span><span class="o">.</span><span class="n">hdfs_path</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">path</span><span class="p">)</span></div>

<div class="viewcode-block" id="TFNodeContext.start_cluster_server"><a class="viewcode-back" href="../../tensorflowonspark.TFSparkNode.html#tensorflowonspark.TFSparkNode.TFNodeContext.start_cluster_server">[docs]</a>  <span class="k">def</span> <span class="nf">start_cluster_server</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_gpus</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">rdma</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Convenience function to access ``TFNode.start_cluster_server`` directly from this object instance.&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">TFNode</span><span class="o">.</span><span class="n">start_cluster_server</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_gpus</span><span class="p">,</span> <span class="n">rdma</span><span class="p">)</span></div>

<div class="viewcode-block" id="TFNodeContext.export_saved_model"><a class="viewcode-back" href="../../tensorflowonspark.TFSparkNode.html#tensorflowonspark.TFSparkNode.TFNodeContext.export_saved_model">[docs]</a>  <span class="k">def</span> <span class="nf">export_saved_model</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sess</span><span class="p">,</span> <span class="n">export_dir</span><span class="p">,</span> <span class="n">tag_set</span><span class="p">,</span> <span class="n">signatures</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Convenience function to access ``TFNode.export_saved_model`` directly from this object instance.&quot;&quot;&quot;</span>
    <span class="n">TFNode</span><span class="o">.</span><span class="n">export_saved_model</span><span class="p">(</span><span class="n">sess</span><span class="p">,</span> <span class="n">export_dir</span><span class="p">,</span> <span class="n">tag_set</span><span class="p">,</span> <span class="n">signatures</span><span class="p">)</span></div>

<div class="viewcode-block" id="TFNodeContext.get_data_feed"><a class="viewcode-back" href="../../tensorflowonspark.TFSparkNode.html#tensorflowonspark.TFSparkNode.TFNodeContext.get_data_feed">[docs]</a>  <span class="k">def</span> <span class="nf">get_data_feed</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">train_mode</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">qname_in</span><span class="o">=</span><span class="s1">&#39;input&#39;</span><span class="p">,</span> <span class="n">qname_out</span><span class="o">=</span><span class="s1">&#39;output&#39;</span><span class="p">,</span> <span class="n">input_mapping</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Convenience function to access ``TFNode.DataFeed`` directly from this object instance.&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">TFNode</span><span class="o">.</span><span class="n">DataFeed</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mgr</span><span class="p">,</span> <span class="n">train_mode</span><span class="p">,</span> <span class="n">qname_in</span><span class="p">,</span> <span class="n">qname_out</span><span class="p">,</span> <span class="n">input_mapping</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="TFSparkNode"><a class="viewcode-back" href="../../tensorflowonspark.TFSparkNode.html#tensorflowonspark.TFSparkNode.TFSparkNode">[docs]</a><span class="k">class</span> <span class="nc">TFSparkNode</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Low-level functions used by the high-level TFCluster APIs to manage cluster state.</span>

<span class="sd">  **This class is not intended for end-users (see TFNode for end-user APIs)**.</span>

<span class="sd">  For cluster management, this wraps the per-node cluster logic as Spark RDD mapPartitions functions, where the RDD is expected to be</span>
<span class="sd">  a &quot;nodeRDD&quot; of the form: ``nodeRDD = sc.parallelize(range(num_executors), num_executors)``.</span>

<span class="sd">  For data feeding, this wraps the feeding logic as Spark RDD mapPartitions functions on a standard &quot;dataRDD&quot;.</span>

<span class="sd">  This also manages a reference to the TFManager &quot;singleton&quot; per executor.  Since Spark can spawn more than one python-worker</span>
<span class="sd">  per executor, this will reconnect to the &quot;singleton&quot; instance as needed.</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="n">mgr</span> <span class="o">=</span> <span class="kc">None</span>                <span class="c1">#: TFManager instance</span>
  <span class="n">cluster_id</span> <span class="o">=</span> <span class="kc">None</span>         <span class="c1">#: Unique ID for a given TensorFlowOnSpark cluster, used for invalidating state for new clusters.</span></div>


<span class="k">def</span> <span class="nf">_get_manager</span><span class="p">(</span><span class="n">cluster_info</span><span class="p">,</span> <span class="n">host</span><span class="p">,</span> <span class="n">executor_id</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Returns this executor&#39;s &quot;singleton&quot; instance of the multiprocessing.Manager, reconnecting per python-worker if needed.</span>

<span class="sd">  Args:</span>
<span class="sd">    :cluster_info: cluster node reservations</span>
<span class="sd">    :host: host IP address</span>
<span class="sd">    :executor_id: unique id per executor (created during initial call to run())</span>

<span class="sd">  Returns:</span>
<span class="sd">    TFManager instance for this executor/python-worker</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">cluster_info</span><span class="p">:</span>
    <span class="k">if</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;host&#39;</span><span class="p">]</span> <span class="o">==</span> <span class="n">host</span> <span class="ow">and</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;executor_id&#39;</span><span class="p">]</span> <span class="o">==</span> <span class="n">executor_id</span><span class="p">:</span>
      <span class="n">addr</span> <span class="o">=</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;addr&#39;</span><span class="p">]</span>
      <span class="n">authkey</span> <span class="o">=</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;authkey&#39;</span><span class="p">]</span>
      <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span> <span class="o">=</span> <span class="n">TFManager</span><span class="o">.</span><span class="n">connect</span><span class="p">(</span><span class="n">addr</span><span class="p">,</span> <span class="n">authkey</span><span class="p">)</span>
      <span class="k">break</span>

  <span class="k">if</span> <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
    <span class="n">msg</span> <span class="o">=</span> <span class="s2">&quot;No TFManager found on this node, please ensure that:</span><span class="se">\n</span><span class="s2">&quot;</span> <span class="o">+</span> \
          <span class="s2">&quot;1. Spark num_executors matches TensorFlow cluster_size</span><span class="se">\n</span><span class="s2">&quot;</span> <span class="o">+</span> \
          <span class="s2">&quot;2. Spark cores/tasks per executor is 1.</span><span class="se">\n</span><span class="s2">&quot;</span> <span class="o">+</span> \
          <span class="s2">&quot;3. Spark dynamic allocation is disabled.&quot;</span>
    <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="n">msg</span><span class="p">)</span>

  <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Connected to TFSparkNode.mgr on </span><span class="si">{0}</span><span class="s2">, executor=</span><span class="si">{1}</span><span class="s2">, state=</span><span class="si">{2}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">host</span><span class="p">,</span> <span class="n">executor_id</span><span class="p">,</span> <span class="nb">str</span><span class="p">(</span><span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;state&#39;</span><span class="p">))))</span>
  <span class="k">return</span> <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span>


<div class="viewcode-block" id="run"><a class="viewcode-back" href="../../tensorflowonspark.TFSparkNode.html#tensorflowonspark.TFSparkNode.run">[docs]</a><span class="k">def</span> <span class="nf">run</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">tf_args</span><span class="p">,</span> <span class="n">cluster_meta</span><span class="p">,</span> <span class="n">tensorboard</span><span class="p">,</span> <span class="n">log_dir</span><span class="p">,</span> <span class="n">queues</span><span class="p">,</span> <span class="n">background</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Wraps the user-provided TensorFlow main function in a Spark mapPartitions function.</span>

<span class="sd">  Args:</span>
<span class="sd">    :fn: TensorFlow &quot;main&quot; function provided by the user.</span>
<span class="sd">    :tf_args: ``argparse`` args, or command line ``ARGV``.  These will be passed to the ``fn``.</span>
<span class="sd">    :cluster_meta: dictionary of cluster metadata (e.g. cluster_id, reservation.Server address, etc).</span>
<span class="sd">    :tensorboard: boolean indicating if the chief worker should spawn a Tensorboard server.</span>
<span class="sd">    :log_dir: directory to save tensorboard event logs.  If None, defaults to a fixed path on local filesystem.</span>
<span class="sd">    :queues: *INTERNAL_USE*</span>
<span class="sd">    :background: boolean indicating if the TensorFlow &quot;main&quot; function should be run in a background process.</span>

<span class="sd">  Returns:</span>
<span class="sd">    A nodeRDD.mapPartitions() function.</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="k">def</span> <span class="nf">_mapfn</span><span class="p">(</span><span class="nb">iter</span><span class="p">):</span>
    <span class="kn">import</span> <span class="nn">tensorflow</span> <span class="k">as</span> <span class="nn">tf</span>

    <span class="c1"># Note: consuming the input iterator helps Pyspark re-use this worker,</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">iter</span><span class="p">:</span>
      <span class="n">executor_id</span> <span class="o">=</span> <span class="n">i</span>

    <span class="c1"># run quick check of GPU infrastructure if using tensorflow-gpu</span>
    <span class="k">if</span> <span class="n">tf</span><span class="o">.</span><span class="n">test</span><span class="o">.</span><span class="n">is_built_with_cuda</span><span class="p">():</span>
      <span class="n">gpus_to_use</span> <span class="o">=</span> <span class="n">gpu_info</span><span class="o">.</span><span class="n">get_gpus</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>

    <span class="c1"># assign TF job/task based on provided cluster_spec template (or use default/null values)</span>
    <span class="n">job_name</span> <span class="o">=</span> <span class="s1">&#39;default&#39;</span>
    <span class="n">task_index</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span>
    <span class="n">cluster_id</span> <span class="o">=</span> <span class="n">cluster_meta</span><span class="p">[</span><span class="s1">&#39;id&#39;</span><span class="p">]</span>
    <span class="n">cluster_template</span> <span class="o">=</span> <span class="n">cluster_meta</span><span class="p">[</span><span class="s1">&#39;cluster_template&#39;</span><span class="p">]</span>
    <span class="k">for</span> <span class="n">jobtype</span> <span class="ow">in</span> <span class="n">cluster_template</span><span class="p">:</span>
      <span class="n">nodes</span> <span class="o">=</span> <span class="n">cluster_template</span><span class="p">[</span><span class="n">jobtype</span><span class="p">]</span>
      <span class="k">if</span> <span class="n">executor_id</span> <span class="ow">in</span> <span class="n">nodes</span><span class="p">:</span>
        <span class="n">job_name</span> <span class="o">=</span> <span class="n">jobtype</span>
        <span class="n">task_index</span> <span class="o">=</span> <span class="n">nodes</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">executor_id</span><span class="p">)</span>
        <span class="k">break</span>

    <span class="c1"># get unique key (hostname, executor_id) for this executor</span>
    <span class="n">host</span> <span class="o">=</span> <span class="n">util</span><span class="o">.</span><span class="n">get_ip_address</span><span class="p">()</span>
    <span class="n">util</span><span class="o">.</span><span class="n">write_executor_id</span><span class="p">(</span><span class="n">executor_id</span><span class="p">)</span>
    <span class="n">port</span> <span class="o">=</span> <span class="mi">0</span>

    <span class="c1"># check for existing TFManagers</span>
    <span class="k">if</span> <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">str</span><span class="p">(</span><span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;state&#39;</span><span class="p">))</span> <span class="o">!=</span> <span class="s2">&quot;&#39;stopped&#39;&quot;</span><span class="p">:</span>
      <span class="k">if</span> <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">cluster_id</span> <span class="o">==</span> <span class="n">cluster_id</span><span class="p">:</span>
        <span class="c1"># raise an exception to force Spark to retry this &quot;reservation&quot; task on another executor</span>
        <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;TFManager already started on </span><span class="si">{0}</span><span class="s2">, executor=</span><span class="si">{1}</span><span class="s2">, state=</span><span class="si">{2}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">host</span><span class="p">,</span> <span class="n">executor_id</span><span class="p">,</span> <span class="nb">str</span><span class="p">(</span><span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;state&quot;</span><span class="p">))))</span>
      <span class="k">else</span><span class="p">:</span>
        <span class="c1"># old state, just continue with creating new manager</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;Ignoring old TFManager with cluster_id </span><span class="si">{0}</span><span class="s2">, requested cluster_id </span><span class="si">{1}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">TFSparkNode</span><span class="o">.</span><span class="n">cluster_id</span><span class="p">,</span> <span class="n">cluster_id</span><span class="p">))</span>

    <span class="c1"># start a TFManager and get a free port</span>
    <span class="c1"># use a random uuid as the authkey</span>
    <span class="n">authkey</span> <span class="o">=</span> <span class="n">uuid</span><span class="o">.</span><span class="n">uuid4</span><span class="p">()</span><span class="o">.</span><span class="n">bytes</span>
    <span class="n">addr</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">if</span> <span class="n">job_name</span> <span class="o">==</span> <span class="s1">&#39;ps&#39;</span><span class="p">:</span>
      <span class="c1"># PS nodes must be remotely accessible in order to shutdown from Spark driver.</span>
      <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span> <span class="o">=</span> <span class="n">TFManager</span><span class="o">.</span><span class="n">start</span><span class="p">(</span><span class="n">authkey</span><span class="p">,</span> <span class="p">[</span><span class="s1">&#39;control&#39;</span><span class="p">,</span> <span class="s1">&#39;error&#39;</span><span class="p">],</span> <span class="s1">&#39;remote&#39;</span><span class="p">)</span>
      <span class="n">addr</span> <span class="o">=</span> <span class="p">(</span><span class="n">host</span><span class="p">,</span> <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">address</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
    <span class="k">else</span><span class="p">:</span>
      <span class="c1"># worker nodes only need to be locally accessible within the executor for data feeding</span>
      <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span> <span class="o">=</span> <span class="n">TFManager</span><span class="o">.</span><span class="n">start</span><span class="p">(</span><span class="n">authkey</span><span class="p">,</span> <span class="n">queues</span><span class="p">)</span>
      <span class="n">addr</span> <span class="o">=</span> <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">address</span>

    <span class="c1"># initialize mgr state</span>
    <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s1">&#39;state&#39;</span><span class="p">,</span> <span class="s1">&#39;running&#39;</span><span class="p">)</span>
    <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">cluster_id</span> <span class="o">=</span> <span class="n">cluster_id</span>

    <span class="c1"># expand Hadoop classpath wildcards for JNI (Spark 2.x)</span>
    <span class="k">if</span> <span class="s1">&#39;HADOOP_PREFIX&#39;</span> <span class="ow">in</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">:</span>
      <span class="n">classpath</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;CLASSPATH&#39;</span><span class="p">]</span>
      <span class="n">hadoop_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;HADOOP_PREFIX&#39;</span><span class="p">],</span> <span class="s1">&#39;bin&#39;</span><span class="p">,</span> <span class="s1">&#39;hadoop&#39;</span><span class="p">)</span>
      <span class="n">hadoop_classpath</span> <span class="o">=</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">check_output</span><span class="p">([</span><span class="n">hadoop_path</span><span class="p">,</span> <span class="s1">&#39;classpath&#39;</span><span class="p">,</span> <span class="s1">&#39;--glob&#39;</span><span class="p">])</span><span class="o">.</span><span class="n">decode</span><span class="p">()</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;CLASSPATH: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">hadoop_classpath</span><span class="p">))</span>
      <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;CLASSPATH&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">classpath</span> <span class="o">+</span> <span class="n">os</span><span class="o">.</span><span class="n">pathsep</span> <span class="o">+</span> <span class="n">hadoop_classpath</span>

    <span class="c1"># start TensorBoard if requested</span>
    <span class="n">tb_pid</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="n">tb_port</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">if</span> <span class="n">tensorboard</span> <span class="ow">and</span> <span class="n">job_name</span> <span class="o">==</span> <span class="s1">&#39;worker&#39;</span> <span class="ow">and</span> <span class="n">task_index</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
      <span class="n">tb_sock</span> <span class="o">=</span> <span class="n">socket</span><span class="o">.</span><span class="n">socket</span><span class="p">(</span><span class="n">socket</span><span class="o">.</span><span class="n">AF_INET</span><span class="p">,</span> <span class="n">socket</span><span class="o">.</span><span class="n">SOCK_STREAM</span><span class="p">)</span>
      <span class="n">tb_sock</span><span class="o">.</span><span class="n">bind</span><span class="p">((</span><span class="s1">&#39;&#39;</span><span class="p">,</span> <span class="mi">0</span><span class="p">))</span>
      <span class="n">tb_port</span> <span class="o">=</span> <span class="n">tb_sock</span><span class="o">.</span><span class="n">getsockname</span><span class="p">()[</span><span class="mi">1</span><span class="p">]</span>
      <span class="n">tb_sock</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
      <span class="n">logdir</span> <span class="o">=</span> <span class="n">log_dir</span> <span class="k">if</span> <span class="n">log_dir</span> <span class="k">else</span> <span class="s2">&quot;tensorboard_</span><span class="si">%d</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="n">executor_id</span>

      <span class="c1"># search for tensorboard in python/bin, PATH, and PYTHONPATH</span>
      <span class="n">pypath</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">executable</span>
      <span class="n">pydir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">pypath</span><span class="p">)</span>
      <span class="n">search_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">pathsep</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="n">pydir</span><span class="p">,</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;PATH&#39;</span><span class="p">],</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;PYTHONPATH&#39;</span><span class="p">]])</span>
      <span class="n">tb_path</span> <span class="o">=</span> <span class="n">util</span><span class="o">.</span><span class="n">find_in_path</span><span class="p">(</span><span class="n">search_path</span><span class="p">,</span> <span class="s1">&#39;tensorboard&#39;</span><span class="p">)</span>                             <span class="c1"># executable in PATH</span>
      <span class="k">if</span> <span class="ow">not</span> <span class="n">tb_path</span><span class="p">:</span>
        <span class="n">tb_path</span> <span class="o">=</span> <span class="n">util</span><span class="o">.</span><span class="n">find_in_path</span><span class="p">(</span><span class="n">search_path</span><span class="p">,</span> <span class="s1">&#39;tensorboard/main.py&#39;</span><span class="p">)</span>                   <span class="c1"># TF 1.3+</span>
      <span class="k">if</span> <span class="ow">not</span> <span class="n">tb_path</span><span class="p">:</span>
        <span class="n">tb_path</span> <span class="o">=</span> <span class="n">util</span><span class="o">.</span><span class="n">find_in_path</span><span class="p">(</span><span class="n">search_path</span><span class="p">,</span> <span class="s1">&#39;tensorflow/tensorboard/__main__.py&#39;</span><span class="p">)</span>    <span class="c1"># TF 1.2-</span>
      <span class="k">if</span> <span class="ow">not</span> <span class="n">tb_path</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Unable to find &#39;tensorboard&#39; in: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">search_path</span><span class="p">))</span>

      <span class="c1"># launch tensorboard</span>
      <span class="n">tb_proc</span> <span class="o">=</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">Popen</span><span class="p">([</span><span class="n">pypath</span><span class="p">,</span> <span class="n">tb_path</span><span class="p">,</span> <span class="s2">&quot;--logdir=</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="n">logdir</span><span class="p">,</span> <span class="s2">&quot;--port=</span><span class="si">%d</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="n">tb_port</span><span class="p">],</span> <span class="n">env</span><span class="o">=</span><span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">)</span>
      <span class="n">tb_pid</span> <span class="o">=</span> <span class="n">tb_proc</span><span class="o">.</span><span class="n">pid</span>

    <span class="c1"># check server to see if this task is being retried (i.e. already reserved)</span>
    <span class="n">client</span> <span class="o">=</span> <span class="n">reservation</span><span class="o">.</span><span class="n">Client</span><span class="p">(</span><span class="n">cluster_meta</span><span class="p">[</span><span class="s1">&#39;server_addr&#39;</span><span class="p">])</span>
    <span class="n">cluster_info</span> <span class="o">=</span> <span class="n">client</span><span class="o">.</span><span class="n">get_reservations</span><span class="p">()</span>
    <span class="n">tmp_sock</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="n">node_meta</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">cluster_info</span><span class="p">:</span>
      <span class="p">(</span><span class="n">nhost</span><span class="p">,</span> <span class="n">nexec</span><span class="p">)</span> <span class="o">=</span> <span class="p">(</span><span class="n">node</span><span class="p">[</span><span class="s1">&#39;host&#39;</span><span class="p">],</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;executor_id&#39;</span><span class="p">])</span>
      <span class="k">if</span> <span class="n">nhost</span> <span class="o">==</span> <span class="n">host</span> <span class="ow">and</span> <span class="n">nexec</span> <span class="o">==</span> <span class="n">executor_id</span><span class="p">:</span>
        <span class="n">node_meta</span> <span class="o">=</span> <span class="n">node</span>
        <span class="n">port</span> <span class="o">=</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;port&#39;</span><span class="p">]</span>

    <span class="c1"># if not already done, register everything we need to set up the cluster</span>
    <span class="k">if</span> <span class="n">node_meta</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
      <span class="c1"># first, find a free port for TF</span>
      <span class="n">tmp_sock</span> <span class="o">=</span> <span class="n">socket</span><span class="o">.</span><span class="n">socket</span><span class="p">(</span><span class="n">socket</span><span class="o">.</span><span class="n">AF_INET</span><span class="p">,</span> <span class="n">socket</span><span class="o">.</span><span class="n">SOCK_STREAM</span><span class="p">)</span>
      <span class="n">tmp_sock</span><span class="o">.</span><span class="n">setsockopt</span><span class="p">(</span><span class="n">socket</span><span class="o">.</span><span class="n">SOL_SOCKET</span><span class="p">,</span> <span class="n">socket</span><span class="o">.</span><span class="n">SO_REUSEADDR</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
      <span class="n">tmp_sock</span><span class="o">.</span><span class="n">bind</span><span class="p">((</span><span class="s1">&#39;&#39;</span><span class="p">,</span> <span class="n">port</span><span class="p">))</span>
      <span class="n">port</span> <span class="o">=</span> <span class="n">tmp_sock</span><span class="o">.</span><span class="n">getsockname</span><span class="p">()[</span><span class="mi">1</span><span class="p">]</span>

      <span class="n">node_meta</span> <span class="o">=</span> <span class="p">{</span>
          <span class="s1">&#39;executor_id&#39;</span><span class="p">:</span> <span class="n">executor_id</span><span class="p">,</span>
          <span class="s1">&#39;host&#39;</span><span class="p">:</span> <span class="n">host</span><span class="p">,</span>
          <span class="s1">&#39;job_name&#39;</span><span class="p">:</span> <span class="n">job_name</span><span class="p">,</span>
          <span class="s1">&#39;task_index&#39;</span><span class="p">:</span> <span class="n">task_index</span><span class="p">,</span>
          <span class="s1">&#39;port&#39;</span><span class="p">:</span> <span class="n">port</span><span class="p">,</span>
          <span class="s1">&#39;tb_pid&#39;</span><span class="p">:</span> <span class="n">tb_pid</span><span class="p">,</span>
          <span class="s1">&#39;tb_port&#39;</span><span class="p">:</span> <span class="n">tb_port</span><span class="p">,</span>
          <span class="s1">&#39;addr&#39;</span><span class="p">:</span> <span class="n">addr</span><span class="p">,</span>
          <span class="s1">&#39;authkey&#39;</span><span class="p">:</span> <span class="n">authkey</span>
      <span class="p">}</span>
      <span class="c1"># register node metadata with server</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;TFSparkNode.reserve: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">node_meta</span><span class="p">))</span>
      <span class="n">client</span><span class="o">.</span><span class="n">register</span><span class="p">(</span><span class="n">node_meta</span><span class="p">)</span>
      <span class="c1"># wait for other nodes to finish reservations</span>
      <span class="n">cluster_info</span> <span class="o">=</span> <span class="n">client</span><span class="o">.</span><span class="n">await_reservations</span><span class="p">()</span>
      <span class="n">client</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>

    <span class="c1"># construct a TensorFlow clusterspec from cluster_info</span>
    <span class="n">sorted_cluster_info</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">cluster_info</span><span class="p">,</span> <span class="n">key</span><span class="o">=</span><span class="k">lambda</span> <span class="n">k</span><span class="p">:</span> <span class="n">k</span><span class="p">[</span><span class="s1">&#39;executor_id&#39;</span><span class="p">])</span>
    <span class="n">spec</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="n">last_executor_id</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span>
    <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">sorted_cluster_info</span><span class="p">:</span>
      <span class="k">if</span> <span class="p">(</span><span class="n">node</span><span class="p">[</span><span class="s1">&#39;executor_id&#39;</span><span class="p">]</span> <span class="o">==</span> <span class="n">last_executor_id</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Duplicate worker/task in cluster_info&quot;</span><span class="p">)</span>
      <span class="n">last_executor_id</span> <span class="o">=</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;executor_id&#39;</span><span class="p">]</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;node: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">node</span><span class="p">))</span>
      <span class="p">(</span><span class="n">njob</span><span class="p">,</span> <span class="n">nhost</span><span class="p">,</span> <span class="n">nport</span><span class="p">)</span> <span class="o">=</span> <span class="p">(</span><span class="n">node</span><span class="p">[</span><span class="s1">&#39;job_name&#39;</span><span class="p">],</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;host&#39;</span><span class="p">],</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;port&#39;</span><span class="p">])</span>
      <span class="n">hosts</span> <span class="o">=</span> <span class="p">[]</span> <span class="k">if</span> <span class="n">njob</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">spec</span> <span class="k">else</span> <span class="n">spec</span><span class="p">[</span><span class="n">njob</span><span class="p">]</span>
      <span class="n">hosts</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">{0}</span><span class="s2">:</span><span class="si">{1}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">nhost</span><span class="p">,</span> <span class="n">nport</span><span class="p">))</span>
      <span class="n">spec</span><span class="p">[</span><span class="n">njob</span><span class="p">]</span> <span class="o">=</span> <span class="n">hosts</span>

    <span class="c1"># update TF_CONFIG if cluster spec has a &#39;master&#39; node (i.e. tf.estimator)</span>
    <span class="k">if</span> <span class="s1">&#39;master&#39;</span> <span class="ow">in</span> <span class="n">spec</span><span class="p">:</span>
      <span class="n">tf_config</span> <span class="o">=</span> <span class="n">json</span><span class="o">.</span><span class="n">dumps</span><span class="p">({</span>
        <span class="s1">&#39;cluster&#39;</span><span class="p">:</span> <span class="n">spec</span><span class="p">,</span>
        <span class="s1">&#39;task&#39;</span><span class="p">:</span> <span class="p">{</span><span class="s1">&#39;type&#39;</span><span class="p">:</span> <span class="n">job_name</span><span class="p">,</span> <span class="s1">&#39;index&#39;</span><span class="p">:</span> <span class="n">task_index</span><span class="p">},</span>
        <span class="s1">&#39;environment&#39;</span><span class="p">:</span> <span class="s1">&#39;cloud&#39;</span>
      <span class="p">})</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;export TF_CONFIG: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">tf_config</span><span class="p">))</span>
      <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;TF_CONFIG&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">tf_config</span>

    <span class="c1"># reserve GPU</span>
    <span class="k">if</span> <span class="n">tf</span><span class="o">.</span><span class="n">test</span><span class="o">.</span><span class="n">is_built_with_cuda</span><span class="p">():</span>
      <span class="n">num_gpus</span> <span class="o">=</span> <span class="n">tf_args</span><span class="o">.</span><span class="n">num_gpus</span> <span class="k">if</span> <span class="s1">&#39;num_gpus&#39;</span> <span class="ow">in</span> <span class="n">tf_args</span> <span class="k">else</span> <span class="mi">1</span>
      <span class="n">gpus_to_use</span> <span class="o">=</span> <span class="n">gpu_info</span><span class="o">.</span><span class="n">get_gpus</span><span class="p">(</span><span class="n">num_gpus</span><span class="p">)</span>
      <span class="n">gpu_str</span> <span class="o">=</span> <span class="s2">&quot;GPUs&quot;</span> <span class="k">if</span> <span class="n">num_gpus</span> <span class="o">&gt;</span> <span class="mi">1</span> <span class="k">else</span> <span class="s2">&quot;GPU&quot;</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;Requested </span><span class="si">{}</span><span class="s2"> </span><span class="si">{}</span><span class="s2">, setting CUDA_VISIBLE_DEVICES=</span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">num_gpus</span><span class="p">,</span> <span class="n">gpu_str</span><span class="p">,</span> <span class="n">gpus_to_use</span><span class="p">))</span>
      <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;CUDA_VISIBLE_DEVICES&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">gpus_to_use</span>

    <span class="c1"># create a context object to hold metadata for TF</span>
    <span class="n">ctx</span> <span class="o">=</span> <span class="n">TFNodeContext</span><span class="p">(</span><span class="n">executor_id</span><span class="p">,</span> <span class="n">job_name</span><span class="p">,</span> <span class="n">task_index</span><span class="p">,</span> <span class="n">spec</span><span class="p">,</span> <span class="n">cluster_meta</span><span class="p">[</span><span class="s1">&#39;default_fs&#39;</span><span class="p">],</span> <span class="n">cluster_meta</span><span class="p">[</span><span class="s1">&#39;working_dir&#39;</span><span class="p">],</span> <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="p">)</span>

    <span class="c1"># release port reserved for TF as late as possible</span>
    <span class="k">if</span> <span class="n">tmp_sock</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
      <span class="n">tmp_sock</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>

    <span class="c1"># Background mode relies reuse of python worker in Spark.</span>
    <span class="k">if</span> <span class="n">background</span><span class="p">:</span>
      <span class="c1"># However, reuse of python worker can&#39;t work on Windows, we need to check if the current</span>
      <span class="c1"># script runs on Windows or not.</span>
      <span class="k">if</span> <span class="n">os</span><span class="o">.</span><span class="n">name</span> <span class="o">==</span> <span class="s1">&#39;nt&#39;</span> <span class="ow">or</span> <span class="n">platform</span><span class="o">.</span><span class="n">system</span><span class="p">()</span> <span class="o">==</span> <span class="s1">&#39;Windows&#39;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Background mode is not supported on Windows.&quot;</span><span class="p">)</span>
      <span class="c1"># Check if the config of reuse python worker is enabled on Spark.</span>
      <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;SPARK_REUSE_WORKER&quot;</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Background mode relies reuse of python worker on Spark. This config &#39;spark.python.worker.reuse&#39; is not enabled on Spark. Please enable it before using background.&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">wrapper_fn</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
      <span class="sd">&quot;&quot;&quot;Wrapper function that sets the sys.argv of the executor.&quot;&quot;&quot;</span>
      <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
        <span class="n">sys</span><span class="o">.</span><span class="n">argv</span> <span class="o">=</span> <span class="n">args</span>
      <span class="n">fn</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="n">context</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">wrapper_fn_background</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
      <span class="sd">&quot;&quot;&quot;Wrapper function that signals exceptions to foreground process.&quot;&quot;&quot;</span>
      <span class="n">errq</span> <span class="o">=</span> <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="s1">&#39;error&#39;</span><span class="p">)</span>
      <span class="k">try</span><span class="p">:</span>
        <span class="n">wrapper_fn</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="n">context</span><span class="p">)</span>
      <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
        <span class="n">errq</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="n">traceback</span><span class="o">.</span><span class="n">format_exc</span><span class="p">())</span>
        <span class="n">errq</span><span class="o">.</span><span class="n">join</span><span class="p">()</span>

    <span class="k">if</span> <span class="n">job_name</span> <span class="o">==</span> <span class="s1">&#39;ps&#39;</span> <span class="ow">or</span> <span class="n">background</span><span class="p">:</span>
      <span class="c1"># invoke the TensorFlow main function in a background thread</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Starting TensorFlow </span><span class="si">{0}</span><span class="s2">:</span><span class="si">{1}</span><span class="s2"> as </span><span class="si">{2}</span><span class="s2"> on cluster node </span><span class="si">{3}</span><span class="s2"> on background process&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
        <span class="n">job_name</span><span class="p">,</span> <span class="n">task_index</span><span class="p">,</span> <span class="n">job_name</span><span class="p">,</span> <span class="n">executor_id</span><span class="p">))</span>

      <span class="n">p</span> <span class="o">=</span> <span class="n">multiprocessing</span><span class="o">.</span><span class="n">Process</span><span class="p">(</span><span class="n">target</span><span class="o">=</span><span class="n">wrapper_fn_background</span><span class="p">,</span> <span class="n">args</span><span class="o">=</span><span class="p">(</span><span class="n">tf_args</span><span class="p">,</span> <span class="n">ctx</span><span class="p">))</span>
      <span class="k">if</span> <span class="n">job_name</span> <span class="o">==</span> <span class="s1">&#39;ps&#39;</span><span class="p">:</span>
        <span class="n">p</span><span class="o">.</span><span class="n">daemon</span> <span class="o">=</span> <span class="kc">True</span>
      <span class="n">p</span><span class="o">.</span><span class="n">start</span><span class="p">()</span>

      <span class="c1"># for ps nodes only, wait indefinitely in foreground thread for a &quot;control&quot; event (None == &quot;stop&quot;)</span>
      <span class="k">if</span> <span class="n">job_name</span> <span class="o">==</span> <span class="s1">&#39;ps&#39;</span><span class="p">:</span>
        <span class="n">queue</span> <span class="o">=</span> <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="s1">&#39;control&#39;</span><span class="p">)</span>
        <span class="n">equeue</span> <span class="o">=</span> <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="s1">&#39;error&#39;</span><span class="p">)</span>
        <span class="n">done</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="k">while</span> <span class="ow">not</span> <span class="n">done</span><span class="p">:</span>
          <span class="k">while</span> <span class="p">(</span><span class="n">queue</span><span class="o">.</span><span class="n">empty</span><span class="p">()</span> <span class="ow">and</span> <span class="n">equeue</span><span class="o">.</span><span class="n">empty</span><span class="p">()):</span>
            <span class="n">time</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
          <span class="k">if</span> <span class="p">(</span><span class="ow">not</span> <span class="n">equeue</span><span class="o">.</span><span class="n">empty</span><span class="p">()):</span>
            <span class="n">e_str</span> <span class="o">=</span> <span class="n">equeue</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>
            <span class="n">equeue</span><span class="o">.</span><span class="n">task_done</span><span class="p">()</span>
            <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;exception in ps:</span><span class="se">\n</span><span class="s2">&quot;</span> <span class="o">+</span> <span class="n">e_str</span><span class="p">)</span>
          <span class="n">msg</span> <span class="o">=</span> <span class="n">queue</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">block</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
          <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Got msg: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">msg</span><span class="p">))</span>
          <span class="k">if</span> <span class="n">msg</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Terminating PS&quot;</span><span class="p">)</span>
            <span class="n">TFSparkNode</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s1">&#39;state&#39;</span><span class="p">,</span> <span class="s1">&#39;stopped&#39;</span><span class="p">)</span>
            <span class="n">done</span> <span class="o">=</span> <span class="kc">True</span>
          <span class="n">queue</span><span class="o">.</span><span class="n">task_done</span><span class="p">()</span>
    <span class="k">else</span><span class="p">:</span>
      <span class="c1"># otherwise, just run TF function in the main executor/worker thread</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Starting TensorFlow </span><span class="si">{0}</span><span class="s2">:</span><span class="si">{1}</span><span class="s2"> on cluster node </span><span class="si">{2}</span><span class="s2"> on foreground thread&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">job_name</span><span class="p">,</span> <span class="n">task_index</span><span class="p">,</span> <span class="n">executor_id</span><span class="p">))</span>
      <span class="n">wrapper_fn</span><span class="p">(</span><span class="n">tf_args</span><span class="p">,</span> <span class="n">ctx</span><span class="p">)</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Finished TensorFlow </span><span class="si">{0}</span><span class="s2">:</span><span class="si">{1}</span><span class="s2"> on cluster node </span><span class="si">{2}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">job_name</span><span class="p">,</span> <span class="n">task_index</span><span class="p">,</span> <span class="n">executor_id</span><span class="p">))</span>

  <span class="k">return</span> <span class="n">_mapfn</span></div>


<div class="viewcode-block" id="train"><a class="viewcode-back" href="../../tensorflowonspark.TFSparkNode.html#tensorflowonspark.TFSparkNode.train">[docs]</a><span class="k">def</span> <span class="nf">train</span><span class="p">(</span><span class="n">cluster_info</span><span class="p">,</span> <span class="n">cluster_meta</span><span class="p">,</span> <span class="n">feed_timeout</span><span class="o">=</span><span class="mi">600</span><span class="p">,</span> <span class="n">qname</span><span class="o">=</span><span class="s1">&#39;input&#39;</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Feeds Spark partitions into the shared multiprocessing.Queue.</span>

<span class="sd">  Args:</span>
<span class="sd">    :cluster_info: node reservation information for the cluster (e.g. host, executor_id, pid, ports, etc)</span>
<span class="sd">    :cluster_meta: dictionary of cluster metadata (e.g. cluster_id, reservation.Server address, etc)</span>
<span class="sd">    :feed_timeout: number of seconds after which data feeding times out (600 sec default)</span>
<span class="sd">    :qname: *INTERNAL_USE*</span>

<span class="sd">  Returns:</span>
<span class="sd">    A dataRDD.mapPartitions() function</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="k">def</span> <span class="nf">_train</span><span class="p">(</span><span class="nb">iter</span><span class="p">):</span>
    <span class="c1"># get shared queue, reconnecting if necessary</span>
    <span class="n">mgr</span> <span class="o">=</span> <span class="n">_get_manager</span><span class="p">(</span><span class="n">cluster_info</span><span class="p">,</span> <span class="n">util</span><span class="o">.</span><span class="n">get_ip_address</span><span class="p">(),</span> <span class="n">util</span><span class="o">.</span><span class="n">read_executor_id</span><span class="p">())</span>
    <span class="k">try</span><span class="p">:</span>
      <span class="n">queue</span> <span class="o">=</span> <span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="n">qname</span><span class="p">)</span>
      <span class="n">equeue</span> <span class="o">=</span> <span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="s1">&#39;error&#39;</span><span class="p">)</span>
    <span class="k">except</span> <span class="p">(</span><span class="ne">AttributeError</span><span class="p">,</span> <span class="ne">KeyError</span><span class="p">):</span>
      <span class="n">msg</span> <span class="o">=</span> <span class="s2">&quot;Queue &#39;</span><span class="si">{}</span><span class="s2">&#39; not found on this node, check for exceptions on other nodes.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">qname</span><span class="p">)</span>
      <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="n">msg</span><span class="p">)</span>

    <span class="n">state</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">mgr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;state&#39;</span><span class="p">))</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;mgr.state=</span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">state</span><span class="p">))</span>
    <span class="n">terminating</span> <span class="o">=</span> <span class="n">state</span> <span class="o">==</span> <span class="s2">&quot;&#39;terminating&#39;&quot;</span>
    <span class="k">if</span> <span class="n">terminating</span><span class="p">:</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;mgr is terminating, skipping partition&quot;</span><span class="p">)</span>
      <span class="n">count</span> <span class="o">=</span> <span class="mi">0</span>
      <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="nb">iter</span><span class="p">:</span>
        <span class="n">count</span> <span class="o">+=</span> <span class="mi">1</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Skipped </span><span class="si">{0}</span><span class="s2"> items from partition&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">count</span><span class="p">))</span>
    <span class="k">else</span><span class="p">:</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Feeding partition </span><span class="si">{0}</span><span class="s2"> into </span><span class="si">{1}</span><span class="s2"> queue </span><span class="si">{2}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">iter</span><span class="p">,</span> <span class="n">qname</span><span class="p">,</span> <span class="n">queue</span><span class="p">))</span>
      <span class="n">count</span> <span class="o">=</span> <span class="mi">0</span>
      <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="nb">iter</span><span class="p">:</span>
        <span class="n">count</span> <span class="o">+=</span> <span class="mi">1</span>
        <span class="n">queue</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="n">block</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

      <span class="c1"># wait for consumers to finish processing all items in queue before &quot;finishing&quot; this iterator</span>
      <span class="n">joinThr</span> <span class="o">=</span> <span class="n">Thread</span><span class="p">(</span><span class="n">target</span><span class="o">=</span><span class="n">queue</span><span class="o">.</span><span class="n">join</span><span class="p">)</span>
      <span class="n">joinThr</span><span class="o">.</span><span class="n">start</span><span class="p">()</span>
      <span class="n">timeout</span> <span class="o">=</span> <span class="n">feed_timeout</span>
      <span class="k">while</span> <span class="p">(</span><span class="n">joinThr</span><span class="o">.</span><span class="n">isAlive</span><span class="p">()):</span>
        <span class="k">if</span> <span class="p">(</span><span class="ow">not</span> <span class="n">equeue</span><span class="o">.</span><span class="n">empty</span><span class="p">()):</span>
          <span class="n">e_str</span> <span class="o">=</span> <span class="n">equeue</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>
          <span class="n">equeue</span><span class="o">.</span><span class="n">task_done</span><span class="p">()</span>
          <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;exception in worker:</span><span class="se">\n</span><span class="s2">&quot;</span> <span class="o">+</span> <span class="n">e_str</span><span class="p">)</span>
        <span class="n">time</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">timeout</span> <span class="o">-=</span> <span class="mi">1</span>
        <span class="k">if</span> <span class="n">timeout</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
          <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Timeout while feeding partition&quot;</span><span class="p">)</span>

      <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Processed </span><span class="si">{0}</span><span class="s2"> items in partition&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">count</span><span class="p">))</span>

    <span class="c1"># check if TF is terminating feed after this partition</span>
    <span class="n">state</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">mgr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;state&#39;</span><span class="p">))</span>
    <span class="n">terminating</span> <span class="o">=</span> <span class="n">state</span> <span class="o">==</span> <span class="s2">&quot;&#39;terminating&#39;&quot;</span>
    <span class="k">if</span> <span class="n">terminating</span><span class="p">:</span>
      <span class="k">try</span><span class="p">:</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;TFSparkNode: requesting stop&quot;</span><span class="p">)</span>
        <span class="n">client</span> <span class="o">=</span> <span class="n">reservation</span><span class="o">.</span><span class="n">Client</span><span class="p">(</span><span class="n">cluster_meta</span><span class="p">[</span><span class="s1">&#39;server_addr&#39;</span><span class="p">])</span>
        <span class="n">client</span><span class="o">.</span><span class="n">request_stop</span><span class="p">()</span>
        <span class="n">client</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
      <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
        <span class="c1"># ignore any errors while requesting stop</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;Error while requesting stop: </span><span class="si">{0}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">e</span><span class="p">))</span>
    <span class="k">return</span> <span class="p">[</span><span class="n">terminating</span><span class="p">]</span>

  <span class="k">return</span> <span class="n">_train</span></div>


<div class="viewcode-block" id="inference"><a class="viewcode-back" href="../../tensorflowonspark.TFSparkNode.html#tensorflowonspark.TFSparkNode.inference">[docs]</a><span class="k">def</span> <span class="nf">inference</span><span class="p">(</span><span class="n">cluster_info</span><span class="p">,</span> <span class="n">feed_timeout</span><span class="o">=</span><span class="mi">600</span><span class="p">,</span> <span class="n">qname</span><span class="o">=</span><span class="s1">&#39;input&#39;</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Feeds Spark partitions into the shared multiprocessing.Queue and returns inference results.</span>

<span class="sd">  Args:</span>
<span class="sd">    :cluster_info: node reservation information for the cluster (e.g. host, executor_id, pid, ports, etc)</span>
<span class="sd">    :feed_timeout: number of seconds after which data feeding times out (600 sec default)</span>
<span class="sd">    :qname: *INTERNAL_USE*</span>

<span class="sd">  Returns:</span>
<span class="sd">    A dataRDD.mapPartitions() function</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="k">def</span> <span class="nf">_inference</span><span class="p">(</span><span class="nb">iter</span><span class="p">):</span>
    <span class="c1"># get shared queue, reconnecting if necessary</span>
    <span class="n">mgr</span> <span class="o">=</span> <span class="n">_get_manager</span><span class="p">(</span><span class="n">cluster_info</span><span class="p">,</span> <span class="n">util</span><span class="o">.</span><span class="n">get_ip_address</span><span class="p">(),</span> <span class="n">util</span><span class="o">.</span><span class="n">read_executor_id</span><span class="p">())</span>
    <span class="k">try</span><span class="p">:</span>
      <span class="n">queue_in</span> <span class="o">=</span> <span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="n">qname</span><span class="p">)</span>
      <span class="n">equeue</span> <span class="o">=</span> <span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="s1">&#39;error&#39;</span><span class="p">)</span>
    <span class="k">except</span> <span class="p">(</span><span class="ne">AttributeError</span><span class="p">,</span> <span class="ne">KeyError</span><span class="p">):</span>
      <span class="n">msg</span> <span class="o">=</span> <span class="s2">&quot;Queue &#39;</span><span class="si">{}</span><span class="s2">&#39; not found on this node, check for exceptions on other nodes.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">qname</span><span class="p">)</span>
      <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="n">msg</span><span class="p">)</span>

    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Feeding partition </span><span class="si">{0}</span><span class="s2"> into </span><span class="si">{1}</span><span class="s2"> queue </span><span class="si">{2}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">iter</span><span class="p">,</span> <span class="n">qname</span><span class="p">,</span> <span class="n">queue_in</span><span class="p">))</span>
    <span class="n">count</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="nb">iter</span><span class="p">:</span>
      <span class="n">count</span> <span class="o">+=</span> <span class="mi">1</span>
      <span class="n">queue_in</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="n">block</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

    <span class="c1"># signal &quot;end of partition&quot;</span>
    <span class="n">queue_in</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="n">marker</span><span class="o">.</span><span class="n">EndPartition</span><span class="p">())</span>

    <span class="c1"># skip empty partitions</span>
    <span class="k">if</span> <span class="n">count</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
      <span class="k">return</span> <span class="p">[]</span>

    <span class="c1"># wait for consumers to finish processing all items in queue before &quot;finishing&quot; this iterator</span>
    <span class="n">joinThr</span> <span class="o">=</span> <span class="n">Thread</span><span class="p">(</span><span class="n">target</span><span class="o">=</span><span class="n">queue_in</span><span class="o">.</span><span class="n">join</span><span class="p">)</span>
    <span class="n">joinThr</span><span class="o">.</span><span class="n">start</span><span class="p">()</span>
    <span class="n">timeout</span> <span class="o">=</span> <span class="n">feed_timeout</span>
    <span class="k">while</span> <span class="p">(</span><span class="n">joinThr</span><span class="o">.</span><span class="n">isAlive</span><span class="p">()):</span>
      <span class="k">if</span> <span class="p">(</span><span class="ow">not</span> <span class="n">equeue</span><span class="o">.</span><span class="n">empty</span><span class="p">()):</span>
        <span class="n">e_str</span> <span class="o">=</span> <span class="n">equeue</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>
        <span class="n">equeue</span><span class="o">.</span><span class="n">task_done</span><span class="p">()</span>
        <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;exception in worker:</span><span class="se">\n</span><span class="s2">&quot;</span> <span class="o">+</span> <span class="n">e_str</span><span class="p">)</span>
      <span class="n">time</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
      <span class="n">timeout</span> <span class="o">-=</span> <span class="mi">1</span>
      <span class="k">if</span> <span class="n">timeout</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Timeout while feeding partition&quot;</span><span class="p">)</span>

    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Processed </span><span class="si">{0}</span><span class="s2"> items in partition&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">count</span><span class="p">))</span>

    <span class="c1"># read result queue</span>
    <span class="n">results</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="n">queue_out</span> <span class="o">=</span> <span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="s1">&#39;output&#39;</span><span class="p">)</span>
    <span class="k">while</span> <span class="n">count</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
      <span class="n">result</span> <span class="o">=</span> <span class="n">queue_out</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">block</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
      <span class="n">results</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">result</span><span class="p">)</span>
      <span class="n">count</span> <span class="o">-=</span> <span class="mi">1</span>
      <span class="n">queue_out</span><span class="o">.</span><span class="n">task_done</span><span class="p">()</span>

    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Finished processing partition&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">results</span>

  <span class="k">return</span> <span class="n">_inference</span></div>


<div class="viewcode-block" id="shutdown"><a class="viewcode-back" href="../../tensorflowonspark.TFSparkNode.html#tensorflowonspark.TFSparkNode.shutdown">[docs]</a><span class="k">def</span> <span class="nf">shutdown</span><span class="p">(</span><span class="n">cluster_info</span><span class="p">,</span> <span class="n">queues</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;input&#39;</span><span class="p">]):</span>
  <span class="sd">&quot;&quot;&quot;Stops all TensorFlow nodes by feeding ``None`` into the multiprocessing.Queues.</span>

<span class="sd">  Args:</span>
<span class="sd">    :cluster_info: node reservation information for the cluster (e.g. host, executor_id, pid, ports, etc).</span>
<span class="sd">    :queues: *INTERNAL_USE*</span>

<span class="sd">  Returns:</span>
<span class="sd">    A nodeRDD.mapPartitions() function</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="k">def</span> <span class="nf">_shutdown</span><span class="p">(</span><span class="nb">iter</span><span class="p">):</span>
    <span class="n">host</span> <span class="o">=</span> <span class="n">util</span><span class="o">.</span><span class="n">get_ip_address</span><span class="p">()</span>
    <span class="n">executor_id</span> <span class="o">=</span> <span class="n">util</span><span class="o">.</span><span class="n">read_executor_id</span><span class="p">()</span>

    <span class="c1"># reconnect to shared queue</span>
    <span class="n">mgr</span> <span class="o">=</span> <span class="n">_get_manager</span><span class="p">(</span><span class="n">cluster_info</span><span class="p">,</span> <span class="n">host</span><span class="p">,</span> <span class="n">executor_id</span><span class="p">)</span>

    <span class="c1"># send SIGTERM to Tensorboard proc (if running)</span>
    <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">cluster_info</span><span class="p">:</span>
      <span class="k">if</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;host&#39;</span><span class="p">]</span> <span class="o">==</span> <span class="n">host</span> <span class="ow">and</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;executor_id&#39;</span><span class="p">]</span> <span class="o">==</span> <span class="n">executor_id</span><span class="p">:</span>
        <span class="n">tb_pid</span> <span class="o">=</span> <span class="n">node</span><span class="p">[</span><span class="s1">&#39;tb_pid&#39;</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">tb_pid</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
          <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Stopping tensorboard (pid=</span><span class="si">{0}</span><span class="s2">)&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">tb_pid</span><span class="p">))</span>
          <span class="n">subprocess</span><span class="o">.</span><span class="n">Popen</span><span class="p">([</span><span class="s2">&quot;kill&quot;</span><span class="p">,</span> <span class="nb">str</span><span class="p">(</span><span class="n">tb_pid</span><span class="p">)])</span>

    <span class="c1"># terminate any listening queues</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Stopping all queues&quot;</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">q</span> <span class="ow">in</span> <span class="n">queues</span><span class="p">:</span>
      <span class="k">try</span><span class="p">:</span>
        <span class="n">queue</span> <span class="o">=</span> <span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="n">q</span><span class="p">)</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Feeding None into </span><span class="si">{0}</span><span class="s2"> queue&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">q</span><span class="p">))</span>
        <span class="n">queue</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="kc">None</span><span class="p">,</span> <span class="n">block</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
      <span class="k">except</span> <span class="p">(</span><span class="ne">AttributeError</span><span class="p">,</span> <span class="ne">KeyError</span><span class="p">):</span>
        <span class="n">msg</span> <span class="o">=</span> <span class="s2">&quot;Queue &#39;</span><span class="si">{}</span><span class="s2">&#39; not found on this node, check for exceptions on other nodes.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">q</span><span class="p">)</span>
        <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="n">msg</span><span class="p">)</span>

    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Setting mgr.state to &#39;stopped&#39;&quot;</span><span class="p">)</span>
    <span class="n">mgr</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s1">&#39;state&#39;</span><span class="p">,</span> <span class="s1">&#39;stopped&#39;</span><span class="p">)</span>
    <span class="k">return</span> <span class="p">[</span><span class="kc">True</span><span class="p">]</span>

  <span class="k">return</span> <span class="n">_shutdown</span></div>
</pre></div>

          </div>
        </div>
      </div>
      <div class="sphinxsidebar" role="navigation" aria-label="main navigation">
        <div class="sphinxsidebarwrapper">
<div id="searchbox" style="display: none" role="search">
  <h3>Quick search</h3>
    <div class="searchformwrapper">
    <form class="search" action="../../search.html" method="get">
      <input type="text" name="q" />
      <input type="submit" value="Go" />
      <input type="hidden" name="check_keywords" value="yes" />
      <input type="hidden" name="area" value="default" />
    </form>
    </div>
</div>
<script type="text/javascript">$('#searchbox').show(0);</script>
        </div>
      </div>
      <div class="clearer"></div>
    </div>
    <div class="related" role="navigation" aria-label="related navigation">
      <h3>Navigation</h3>
      <ul>
        <li class="right" style="margin-right: 10px">
          <a href="../../genindex.html" title="General Index"
             >index</a></li>
        <li class="right" >
          <a href="../../py-modindex.html" title="Python Module Index"
             >modules</a> |</li>
        <li class="nav-item nav-item-0"><a href="../../index.html">TensorFlowOnSpark 1.3.3 documentation</a> &#187;</li>
          <li class="nav-item nav-item-1"><a href="../index.html" >Module code</a> &#187;</li> 
      </ul>
    </div>
    <div class="footer" role="contentinfo">
        &#169; Copyright 2018, Yahoo Inc.
      Created using <a href="http://sphinx-doc.org/">Sphinx</a> 1.7.9.
    </div>
  </body>
</html>